#!/usr/bin/env python
# coding: utf-8

# 数值方法3：偏微分方程1 使用有限差分法解一维热传导（扩散）方程
# 有热源情况

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.animation as animation
from matplotlib.animation import FFMpegWriter

# 热传导参数
a = 0.082
dx = 0.02
X = 1
x = []
x = np.linspace(0, X, int(X / dx + 1))
dt = 0.0001
T = 0.5
t = []
t = np.linspace(0, T, int(T / dt + 1))

# 算子
A = (-2 * np.eye(len(x), len(x), dtype=int)
     + np.diag(np.diag(np.ones([len(x), len(x)]), k=1), k=1)
     + np.diag(np.diag(np.ones([len(x), len(x)]), k=1), k=-1))

# 初始化
u = 0 * x
# 边界
m1 = 0 + 0.0 * np.sin(t)
m2 = 0 - 0.0 * np.sin(8 * t)
# 热源
list_f = [75 for i in range(len(x))]
f = np.mat(list_f).T
print(f)
# 解方程组
list_u = [37 for i in range(len(x))]
u = np.mat(list_u).T


# print(u.shape)
for n in range(len(t) - 1):
    u = np.append(u, np.add(u[:, n], np.add(a * a / dx / dx * np.mat(A) * u[:, n], f) * dt), axis=1)
    u[0, n + 1] = m1[n + 1]
    u[-1, n + 1] = m2[n + 1]
    print(u)

# 绘制边界线
Y = np.array(u)

X = np.array(x)
plt.plot(X, Y[:, -1])
plt.title('FinalStates Temperature')
plt.savefig("FinalStatesT.png")  # 保存图片
plt.show()
print(X.shape, Y.shape)
# 3D曲面
# get_ipython().run_line_magic('matplotlib', 'notebook')
# 准备数据
x_x, y_y = np.meshgrid(t, x)

# 绘制图片
fig = plt.figure("3D Surface", facecolor="lightgray")
plt.title("Temperature3D", fontsize=18)

# 设置为3D图片类型
ax3d = Axes3D(fig, auto_add_to_figure=False)
fig.add_axes(ax3d)

ax3d.set_xlabel("time")
ax3d.set_ylabel("x")
ax3d.set_zlabel("temperature")
plt.tick_params(labelsize=10)

ax3d.plot_surface(x_x, y_y, u, cstride=20, rstride=20, cmap="jet")

plt.savefig("Temperature3D.png")  # 保存图片
plt.show()

# 初始化画布
fig = plt.figure()
# plt.xlim(0,X)
plt.ylim(0, np.max(u))  # np.max(psi)
plt.grid(ls='--')

ys = u[:, 0]
Figure = plt.plot(x, ys, c='blue', alpha=0.8)[0]


# 更新函数
def updata(num):
    ys = u[:, num]
    Figure.set_data(x, ys)
    return Figure


# animation库绘制动图
ani = animation.FuncAnimation(fig=fig, func=updata, frames=np.arange(0, int(len(t) / 2)),
                              interval=2)  # frames帧数，interval间隔（帧）
plt.show()

## 保存为mp4，运行速度较慢，不保存时注释掉。
# mywriter = FFMpegWriter(fps=60)
# ani.save('TemperatureHistory.MP4',writer=mywriter)